#!/bin/bash

# DevDeploy 分割API测试脚本
# 用法: ./test_segment_api.sh [port] [image_path]
# 默认端口: 8796
# 默认图片: seg.jpg

PORT=${1:-8796}
IMAGE_PATH=${2:-seg.jpg}

echo "🧪 DevDeploy分割API测试脚本"
echo "================================"
echo "测试端口: ${PORT}"
echo "测试图片: ${IMAGE_PATH}"
echo ""

# 检查图片文件是否存在
if [ ! -f "$IMAGE_PATH" ]; then
    echo "❌ 错误: 测试图片不存在: ${IMAGE_PATH}"
    echo "请指定正确的图片路径，或确保 seg.jpg 在当前目录"
    exit 1
fi

# 检查服务是否运行
echo "1. 检查服务状态..."
if curl -s http://localhost:${PORT}/health > /dev/null; then
    echo "✅ 分割API服务运行正常"
else
    echo "❌ API服务未运行（端口 ${PORT}）"
    echo "请先启动服务: ./script/start_cv_api.sh --segment <model_path> ${PORT}"
    exit 1
fi

# 获取配置信息
echo ""
echo "2. 获取服务配置..."
echo "配置信息:"
CONFIG_RESPONSE=$(curl -s http://localhost:${PORT}/config)
echo "$CONFIG_RESPONSE" | python3 -m json.tool

# 检查分割服务是否启用
SEGMENT_ENABLED=$(echo "$CONFIG_RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('segmentation_enabled', False))")
if [ "$SEGMENT_ENABLED" != "True" ]; then
    echo ""
    echo "⚠️  警告: 分割服务未启用"
    echo "请确保使用 --segment 参数启动服务"
    echo "例如: ./script/start_cv_api.sh --segment /path/to/segment_model.onnx"
    exit 1
fi

# 测试分割功能
echo ""
echo "3. 测试分割功能（使用base64编码）..."
echo "使用测试图片进行分割: ${IMAGE_PATH}"

# 将图片转换为base64编码
BASE64_FILE="seg_base64.txt"
echo "正在生成base64编码文件: ${BASE64_FILE}..."

# 兼容Linux和Mac的base64命令
# Linux: base64 -w 0 (无换行)
# Mac: base64 -i (指定输入文件) 或 base64 < file (重定向)
if base64 -w 0 "$IMAGE_PATH" > "${BASE64_FILE}" 2>/dev/null; then
    # Linux版本成功
    :
elif base64 -i "$IMAGE_PATH" > "${BASE64_FILE}" 2>/dev/null; then
    # Mac版本成功，需要去掉换行符
    tr -d '\n' < "${BASE64_FILE}" > "${BASE64_FILE}.tmp" && mv "${BASE64_FILE}.tmp" "${BASE64_FILE}"
elif base64 < "$IMAGE_PATH" > "${BASE64_FILE}" 2>/dev/null; then
    # 使用重定向方式（Mac备用方案），需要去掉换行符
    tr -d '\n' < "${BASE64_FILE}" > "${BASE64_FILE}.tmp" && mv "${BASE64_FILE}.tmp" "${BASE64_FILE}"
else
    echo "❌ 无法生成base64编码，请检查base64命令是否可用"
    exit 1
fi

if [ $? -eq 0 ]; then
    echo "✅ Base64编码文件已生成: ${BASE64_FILE}"
else
    echo "❌ 生成base64编码文件失败"
    exit 1
fi

# 使用Python创建JSON文件，避免shell转义问题
TEMP_JSON=$(mktemp)
python3 << EOF > "${TEMP_JSON}"
import json

with open("${BASE64_FILE}", "r") as f:
    image_base64 = f.read().strip()

data = {"image": image_base64}
print(json.dumps(data))
EOF

# 发送分割请求
echo ""
echo "4. 发送分割请求..."
RESPONSE=$(curl -s -X POST "http://localhost:${PORT}/api/v1/segment/single" \
    -H "Content-Type: application/json" \
    -d "@${TEMP_JSON}")

# 清理临时文件
rm -f "${TEMP_JSON}"

# 检查响应
if [ $? -ne 0 ]; then
    echo "❌ 请求失败"
    exit 1
fi

# 解析响应
echo ""
echo "5. 分割结果:"
echo "$RESPONSE" | python3 -m json.tool

# 检查是否成功
SUCCESS=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('success', False))" 2>/dev/null)
if [ "$SUCCESS" != "True" ]; then
    echo ""
    echo "❌ 分割失败，请检查响应信息"
    exit 1
fi

# 保存返回的分割掩码图像
echo ""
echo "6. 保存分割掩码图像..."
MASK_BASE64=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('segmentation_mask_base64', ''))" 2>/dev/null)

if [ -n "$MASK_BASE64" ]; then
    # 保存base64字符串到txt文件
    MASK_BASE64_FILE="seg_result_mask_base64.txt"
    echo "$MASK_BASE64" > "$MASK_BASE64_FILE"
    echo "✅ 分割掩码base64已保存: ${MASK_BASE64_FILE}"
    
    # 处理base64字符串（移除数据URI前缀）
    CLEAN_BASE64=$(echo "$MASK_BASE64" | sed 's/^data:image\/[^;]*;base64,//')
    
    # 解码并保存PNG图片
    OUTPUT_FILE="seg_result_mask.png"
    echo "$CLEAN_BASE64" | base64 -d > "$OUTPUT_FILE" 2>/dev/null || \
    echo "$CLEAN_BASE64" | python3 -c "import sys, base64; sys.stdout.buffer.write(base64.b64decode(sys.stdin.read()))" > "$OUTPUT_FILE" 2>/dev/null
    
    if [ -f "$OUTPUT_FILE" ] && [ -s "$OUTPUT_FILE" ]; then
        echo "✅ 分割掩码PNG已保存: ${OUTPUT_FILE}"
        
        # 显示图片信息
        if command -v file &> /dev/null; then
            echo "   文件信息: $(file "$OUTPUT_FILE")"
        fi
    else
        echo "⚠️  保存分割掩码PNG失败，但base64已保存"
    fi
else
    echo "⚠️  响应中没有包含分割掩码图像"
fi

# 保存返回的彩色掩码图像
echo ""
echo "7. 保存彩色掩码图像..."
COLOR_MASK_BASE64=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('color_mask_base64', ''))" 2>/dev/null)

if [ -n "$COLOR_MASK_BASE64" ]; then
    # 保存base64字符串到txt文件
    COLOR_MASK_BASE64_FILE="seg_result_color_mask_base64.txt"
    echo "$COLOR_MASK_BASE64" > "$COLOR_MASK_BASE64_FILE"
    echo "✅ 彩色掩码base64已保存: ${COLOR_MASK_BASE64_FILE}"
    
    # 处理base64字符串（移除数据URI前缀）
    CLEAN_COLOR_BASE64=$(echo "$COLOR_MASK_BASE64" | sed 's/^data:image\/[^;]*;base64,//')
    
    # 解码并保存PNG图片
    COLOR_OUTPUT_FILE="seg_result_color_mask.png"
    echo "$CLEAN_COLOR_BASE64" | base64 -d > "$COLOR_OUTPUT_FILE" 2>/dev/null || \
    echo "$CLEAN_COLOR_BASE64" | python3 -c "import sys, base64; sys.stdout.buffer.write(base64.b64decode(sys.stdin.read()))" > "$COLOR_OUTPUT_FILE" 2>/dev/null
    
    if [ -f "$COLOR_OUTPUT_FILE" ] && [ -s "$COLOR_OUTPUT_FILE" ]; then
        echo "✅ 彩色掩码PNG已保存: ${COLOR_OUTPUT_FILE}"
        
        # 显示图片信息
        if command -v file &> /dev/null; then
            echo "   文件信息: $(file "$COLOR_OUTPUT_FILE")"
        fi
    else
        echo "⚠️  保存彩色掩码PNG失败，但base64已保存"
    fi
else
    echo "⚠️  响应中没有包含彩色掩码图像"
fi

# 显示其他信息
echo ""
echo "8. 分割结果详情:"
IMAGE_SHAPE=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('image_shape', 'N/A'))" 2>/dev/null)
NUM_CLASSES=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('num_classes', 'N/A'))" 2>/dev/null)
CLASS_NAMES=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); names=data.get('class_names', []); print(', '.join(str(n) for n in names) if names else 'N/A')" 2>/dev/null)

echo "   图像尺寸: ${IMAGE_SHAPE}"
echo "   类别数量: ${NUM_CLASSES}"
echo "   类别名称: ${CLASS_NAMES}"

# 显示API文档地址
echo ""
echo "9. API文档地址:"
echo "   http://localhost:${PORT}/docs"
echo ""

# 清理临时base64文件（可选）
read -p "是否删除临时base64文件 ${BASE64_FILE}? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    rm -f "${BASE64_FILE}"
    echo "✅ 已删除临时文件"
fi

echo ""
echo "✅ 测试完成！"

